'''
Created on Jul 16, 2009

@author: xin
@author: Mikael Rouson
'''
import numpy
import extend_sym

def wneigh_extract(X, dims):
    """  
    @param X     N1 x N2 x M: data = M images of size N1 x N2, each image is a
                 wavelet sub-band.
    @param dims  M x 1: size of the neighborhood to be extracted (extract 
                 sqrt(dims(i))x sqrt(dims(i))- or 5- neighborhood of sub-band 
                 X(:,:,i)).
    
    @return
         Y: N1N2 x L: each row of Y is a wavelet neighborhood extracted from X.
    """
    # Verify the size of X and the neighborhood:
    X = numpy.array(X)
    M = X.shape[0]
    N1 = X.shape[1]
    N2 = X.shape[2]
    dims = dims.ravel()
    
    if len(dims) != M :
        print "Input contains %d Subbands, neighborhood size does not match"\
        % (M)
    if any(1 - numpy.mod(dims, 2)):
        print "Neighborhoods must have an odd number of elements"
    
    # Define the neighborhood shape
    tmp = numpy.floor(numpy.sqrt(dims)) ** 2
    flag = (dims == 5)
    if any((tmp != dims) * (1 - flag)):
        print"Neighborhoods must have an odd number of elements"
    
    # Fill in the neighborhoods
    dims = numpy.ceil(numpy.sqrt(dims))
    # square neighborhood: flag=0, dims=side. 
    # 5 neighborhood: flag=1, dims=3 (3 square neighborhood for extension).   
    l = 0
    Y = []
    for m in range(M):
        n = numpy.floor(dims[m] / 2)
        Xext = extend_sym.extend_sym(X[m, :, :], n)
        if flag[m]:  # list of neighbors for  5 neighborhood (center in (1,1)):
            liy = numpy.array([0, 1, 1, 1, 2])
            lix = numpy.array([1, 0, 1, 2, 1])
            for i in range(len(liy)):
                iy = liy[i];
                ix = lix[i];
                tmp = Xext[ix:N1 + ix, iy:N2 + iy]
                
                length = len(tmp.ravel())
                Y.append(tmp.ravel())
                l = l + 1
        else:   # list of neighbors for square neighborhood (top
                # left corner in  (0,0)):
            for iy in range(int(dims[m])):
                for ix in range(int(dims[m])):
                    tmp = Xext[ix:N1 + ix, iy:N2 + iy]
                    length = len(tmp.ravel())
                    Y.append(tmp.ravel())
                    l = l + 1
    return Y
